package com.ruoyi.websocket.handler;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.web.util.UriTemplate;

import java.io.IOException;
import java.net.URI;
import java.util.*;
import java.util.concurrent.*;

@Component
public class ChatWebSocketHandle extends TextWebSocketHandler {

    private static final Logger log = LoggerFactory.getLogger(ChatWebSocketHandle.class);

    /* ======================== 常量定义 ======================== */
    private static final String TYPE_CHAT = "CHAT";
    private static final String TYPE_GROUP_CHAT = "GROUP_CHAT";
    private static final String TYPE_BROADCAST = "BROADCAST";
    private static final String TYPE_HEARTBEAT = "HEARTBEAT";

    // 新增：订阅与推送
    private static final String TYPE_PUSH_SUBSCRIBE = "PUSH_SUBSCRIBE";
    private static final String TYPE_PUSH_UNSUBSCRIBE = "PUSH_UNSUBSCRIBE";
    private static final String TYPE_PUSH = "PUSH"; // 依据 topic 推送

    // 新增：多对话（会话）管理
    private static final String TYPE_START_DIALOG = "START_DIALOG";
    private static final String TYPE_END_DIALOG = "END_DIALOG";
    private static final String TYPE_DIALOG_MSG = "DIALOG_MSG";

    private static final Set<String> SUPPORTED_TYPES = new HashSet<>(Arrays.asList(
            TYPE_CHAT, TYPE_GROUP_CHAT, TYPE_BROADCAST, TYPE_HEARTBEAT,
            TYPE_PUSH_SUBSCRIBE, TYPE_PUSH_UNSUBSCRIBE, TYPE_PUSH,
            TYPE_START_DIALOG, TYPE_END_DIALOG, TYPE_DIALOG_MSG
    ));

    /* ======================== 数据结构 ======================== */

    // openId -> 该用户的所有 WebSocketSession
    private final Map<String, List<WebSocketSession>> sessionMap = new ConcurrentHashMap<>();

    // 主题订阅：topic -> 订阅该主题的所有会话
    private final Map<String, Set<WebSocketSession>> topicSessions = new ConcurrentHashMap<>();

    // 反向索引（断开连接时快速清理）：session -> 已订阅的所有主题
    private final Map<WebSocketSession, Set<String>> sessionTopics = new ConcurrentHashMap<>();

    // 对话管理：dialogId -> 参与该对话的会话集合（可跨用户）
    private final Map<String, Set<WebSocketSession>> dialogSessions = new ConcurrentHashMap<>();

    // 反向索引：session -> 其参与的对话集合
    private final Map<WebSocketSession, Set<String>> sessionDialogs = new ConcurrentHashMap<>();

    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

    // 只匹配路径，避免 host/query 干扰
    private static final UriTemplate URI_TEMPLATE = new UriTemplate("/webSocket/chat/{openId}");

    /* ======================== 工具方法 ======================== */

    private String getOpenId(WebSocketSession session) {
        URI uri = session.getUri();
        if (uri == null) return null;
        String path = uri.getPath(); // 只取 path
        Map<String, String> variables = URI_TEMPLATE.match(path);
        return variables.get("openId");
    }

    private static String asText(JsonNode node, String field, String defVal) {
        if (node == null) return defVal;
        JsonNode n = node.path(field);
        return n.isMissingNode() || n.isNull() ? defVal : n.asText(defVal);
    }

    private static List<String> asStringList(JsonNode node, String field) {
        List<String> list = new ArrayList<>();
        JsonNode arr = node.path(field);
        if (arr.isArray()) {
            arr.forEach(x -> list.add(x.asText()));
        }
        return list;
    }

    /* ======================== 生命周期回调 ======================== */

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String openId = getOpenId(session);
        if (openId != null && !openId.trim().isEmpty()) {
            addSessionToUser(openId, session);
            log.info("用户 {} 连接建立，会话ID: {}", openId, session.getId());
            session.sendMessage(new TextMessage(createSystemPayload("CONNECTED", "连接已建立")));
        } else {
            log.warn("连接建立但未提供有效的openId，会话ID: {}", session.getId());
            session.close(CloseStatus.NOT_ACCEPTABLE.withReason("未提供有效用户标识"));
        }
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String openId = getOpenId(session);
        if (openId == null) {
            session.close(CloseStatus.NOT_ACCEPTABLE.withReason("无效会话"));
            return;
        }

        String payload = message.getPayload();
        log.debug("收到来自用户 {} 的消息: {}", openId, payload);

        JsonNode root;
        try {
            root = OBJECT_MAPPER.readTree(payload);
        } catch (Exception e) {
            log.error("JSON 解析失败: {}", e.getMessage());
            session.sendMessage(new TextMessage(createErrorPayload("E_JSON", "消息格式错误")));
            return;
        }

        String type = asText(root, "type", "");
        if (!SUPPORTED_TYPES.contains(type)) {
            log.warn("未知的消息类型: {}", type);
            session.sendMessage(new TextMessage(createErrorPayload("E_TYPE", "未知消息类型")));
            return;
        }

        try {
            switch (type) {
                case TYPE_CHAT: {
                    String content = asText(root, "content", "");
                    String target = asText(root, "target", null);
                    if (target == null || target.isEmpty()) {
                        session.sendMessage(new TextMessage(createErrorPayload("E_ARG", "缺少目标用户")));
                        break;
                    }
                    sendToUser(target, new TextMessage(createMessagePayload(openId, content, TYPE_CHAT)));
                    break;
                }
                case TYPE_GROUP_CHAT: {
                    String content = asText(root, "content", "");
                    String groupId = asText(root, "target", null);
                    if (groupId == null || groupId.isEmpty()) {
                        session.sendMessage(new TextMessage(createErrorPayload("E_ARG", "缺少群组ID")));
                        break;
                    }
                    broadcastToGroup(groupId, new TextMessage(createMessagePayload(openId, content, TYPE_GROUP_CHAT)));
                    break;
                }
                case TYPE_BROADCAST: {
                    String content = asText(root, "content", "");
                    broadcastToAll(new TextMessage(createMessagePayload(openId, content, TYPE_BROADCAST)));
                    break;
                }
                case TYPE_HEARTBEAT: {
                    session.sendMessage(new TextMessage(createSystemPayload("HEARTBEAT", "pong")));
                    break;
                }

                /* ======== 新增：订阅/取消订阅/按主题推送 ======== */
                case TYPE_PUSH_SUBSCRIBE: {
                    List<String> topics = asStringList(root, "topics");
                    if (topics.isEmpty()) {
                        session.sendMessage(new TextMessage(createErrorPayload("E_ARG", "缺少 topics")));
                        break;
                    }
                    int added = 0;
                    for (String topic : topics) {
                        if (topic == null || topic.isEmpty()) continue;
                        topicSessions.computeIfAbsent(topic, k -> ConcurrentHashMap.newKeySet()).add(session);
                        sessionTopics.computeIfAbsent(session, k -> ConcurrentHashMap.newKeySet()).add(topic);
                        added++;
                    }
                    session.sendMessage(new TextMessage(createSystemEvent("PUSH_SUBSCRIBE_OK",
                            "已订阅主题", mapOf("count", added, "topics", topics))));
                    break;
                }
                case TYPE_PUSH_UNSUBSCRIBE: {
                    List<String> topics = asStringList(root, "topics");
                    boolean all = topics.isEmpty(); // 如果未传，默认全部取消
                    int removed = 0;

                    if (all) {
                        Set<String> myTopics = sessionTopics.getOrDefault(session, Collections.emptySet());
                        for (String t : new ArrayList<>(myTopics)) {
                            removeSessionFromTopic(t, session);
                            removed++;
                        }
                    } else {
                        for (String t : topics) {
                            if (t == null || t.isEmpty()) continue;
                            if (removeSessionFromTopic(t, session)) removed++;
                        }
                    }
                    session.sendMessage(new TextMessage(createSystemEvent("PUSH_UNSUBSCRIBE_OK",
                            all ? "已取消全部订阅" : "已取消指定订阅",
                            mapOf("count", removed, "topics", topics))));
                    break;
                }
                case TYPE_PUSH: {
                    // 按主题推送（通常由服务端触发，这里也允许某客户端触发示例）
                    String topic = asText(root, "topic", null);
                    String content = asText(root, "content", "");
                    if (topic == null || topic.isEmpty()) {
                        session.sendMessage(new TextMessage(createErrorPayload("E_ARG", "缺少 topic")));
                        break;
                    }
                    broadcastToTopic(topic, new TextMessage(createMessagePayload(openId, content, TYPE_PUSH)));
                    break;
                }

                /* ======== 新增：对话（会话）管理 ======== */
                case TYPE_START_DIALOG: {
                    // 可由客户端自带 dialogId，否则服务端生成
                    String dialogId = asText(root, "dialogId", "");
                    if (dialogId.isEmpty()) {
                        dialogId = UUID.randomUUID().toString().replace("-", "");
                    }
                    dialogSessions.computeIfAbsent(dialogId, k -> ConcurrentHashMap.newKeySet()).add(session);
                    sessionDialogs.computeIfAbsent(session, k -> ConcurrentHashMap.newKeySet()).add(dialogId);

                    // 可选：邀请其他参与者（用户 openId 列表）
                    List<String> participants = asStringList(root, "participants"); // openIds
                    for (String pid : participants) {
                        List<WebSocketSession> targetSessions = sessionMap.getOrDefault(pid, Collections.emptyList());
                        for (WebSocketSession s : targetSessions) {
                            dialogSessions.get(dialogId).add(s);
                            sessionDialogs.computeIfAbsent(s, k -> ConcurrentHashMap.newKeySet()).add(dialogId);
                        }
                    }
                    // 将对话开始事件广播给对话内所有参与者
                    broadcastToDialog(dialogId, new TextMessage(createSystemEvent("DIALOG_STARTED",
                            "对话已创建", mapOf("dialogId", dialogId, "participants", participants))));
                    break;
                }
                case TYPE_DIALOG_MSG: {
                    String dialogId = asText(root, "dialogId", "");
                    String content = asText(root, "content", "");
                    if (dialogId.isEmpty()) {
                        session.sendMessage(new TextMessage(createErrorPayload("E_ARG", "缺少 dialogId")));
                        break;
                    }
                    if (!dialogSessions.containsKey(dialogId)) {
                        session.sendMessage(new TextMessage(createErrorPayload("E_NOT_FOUND", "对话不存在")));
                        break;
                    }
                    broadcastToDialog(dialogId, new TextMessage(createDialogPayload(dialogId, openId, content)));
                    break;
                }
                case TYPE_END_DIALOG: {
                    String dialogId = asText(root, "dialogId", "");
                    if (dialogId.isEmpty()) {
                        session.sendMessage(new TextMessage(createErrorPayload("E_ARG", "缺少 dialogId")));
                        break;
                    }
                    closeDialog(dialogId, "对话已结束");
                    break;
                }

                default:
                    session.sendMessage(new TextMessage(createErrorPayload("E_TYPE", "未知消息类型")));
            }
        } catch (Exception e) {
            log.error("处理消息时发生错误: {}", e.getMessage(), e);
            session.sendMessage(new TextMessage(createErrorPayload("E_RUNTIME", "服务器内部错误")));
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        String openId = getOpenId(session);
        log.error("用户 {} 的传输错误: {}", openId, exception.getMessage(), exception);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        String openId = getOpenId(session);
        if (openId != null) {
            removeSessionFromUser(openId, session);
            cleanupSessionFromAllTopics(session);
            cleanupSessionFromAllDialogs(session);
            log.info("用户 {} 连接关闭，原因: {}", openId, status.getReason());
        }
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /* ======================== 用户会话组 ======================== */

    private void addSessionToUser(String openId, WebSocketSession session) {
        sessionMap.computeIfAbsent(openId, k -> new CopyOnWriteArrayList<>()).add(session);
        log.debug("会话已添加到用户 {}，当前用户会话数: {}", openId, sessionMap.get(openId).size());
    }

    private void removeSessionFromUser(String openId, WebSocketSession session) {
        List<WebSocketSession> sessions = sessionMap.get(openId);
        if (sessions != null) {
            sessions.remove(session);
            log.debug("会话已从用户 {} 移除，剩余会话数: {}", openId, sessions.size());
            if (sessions.isEmpty()) {
                sessionMap.remove(openId);
                log.debug("用户 {} 的所有会话已移除", openId);
            }
        }
    }

    /* ======================== 主题订阅与推送 ======================== */

    private boolean removeSessionFromTopic(String topic, WebSocketSession session) {
        Set<WebSocketSession> set = topicSessions.get(topic);
        boolean removed = false;
        if (set != null) {
            removed = set.remove(session);
            if (set.isEmpty()) topicSessions.remove(topic);
        }
        Set<String> my = sessionTopics.get(session);
        if (my != null) {
            my.remove(topic);
            if (my.isEmpty()) sessionTopics.remove(session);
        }
        return removed;
    }

    private void cleanupSessionFromAllTopics(WebSocketSession session) {
        Set<String> topics = sessionTopics.getOrDefault(session, Collections.emptySet());
        for (String t : new ArrayList<>(topics)) {
            removeSessionFromTopic(t, session);
        }
    }

    private void broadcastToTopic(String topic, TextMessage message) {
        Set<WebSocketSession> sessions = topicSessions.get(topic);
        if (sessions == null || sessions.isEmpty()) {
            log.debug("主题 {} 没有订阅者", topic);
            return;
        }
        int total = 0, ok = 0;
        for (WebSocketSession s : sessions) {
            total++;
            if (s.isOpen()) {
                try {
                    s.sendMessage(message);
                    ok++;
                } catch (IOException e) {
                    log.error("向主题 {} 发送失败: {}", topic, e.getMessage());
                }
            }
        }
        log.debug("向主题 {} 推送完成，成功: {}/{}", topic, ok, total);
    }

    /* ======================== 多对话管理 ======================== */

    private void broadcastToDialog(String dialogId, TextMessage message) {
        Set<WebSocketSession> sessions = dialogSessions.get(dialogId);
        if (sessions == null || sessions.isEmpty()) {
            log.debug("对话 {} 无参与者", dialogId);
            return;
        }
        int total = 0, ok = 0;
        for (WebSocketSession s : sessions) {
            total++;
            if (s.isOpen()) {
                try {
                    s.sendMessage(message);
                    ok++;
                } catch (IOException e) {
                    log.error("向对话 {} 发送失败: {}", dialogId, e.getMessage());
                }
            }
        }
        log.debug("对话 {} 广播完成，成功: {}/{}", dialogId, ok, total);
    }

    private void closeDialog(String dialogId, String reason) {
        Set<WebSocketSession> sessions = dialogSessions.remove(dialogId);
        if (sessions == null) return;
        TextMessage msg = new TextMessage(createSystemEvent("DIALOG_ENDED", reason, mapOf("dialogId", dialogId)));
        for (WebSocketSession s : sessions) {
            try { s.sendMessage(msg); } catch (IOException ignored) {}
            Set<String> ds = sessionDialogs.get(s);
            if (ds != null) {
                ds.remove(dialogId);
                if (ds.isEmpty()) sessionDialogs.remove(s);
            }
        }
        log.debug("对话 {} 已关闭并清理", dialogId);
    }

    private void cleanupSessionFromAllDialogs(WebSocketSession session) {
        Set<String> dialogs = sessionDialogs.getOrDefault(session, Collections.emptySet());
        for (String d : new ArrayList<>(dialogs)) {
            Set<WebSocketSession> set = dialogSessions.get(d);
            if (set != null) {
                set.remove(session);
                if (set.isEmpty()) {
                    dialogSessions.remove(d);
                    log.debug("对话 {} 无参与者，已回收", d);
                }
            }
        }
        sessionDialogs.remove(session);
    }

    /* ======================== 广播/用户发送 ======================== */

    private void sendToUser(String openId, TextMessage message) {
        List<WebSocketSession> sessions = sessionMap.get(openId);
        if (sessions != null && !sessions.isEmpty()) {
            for (WebSocketSession s : sessions) {
                if (s.isOpen()) {
                    try {
                        s.sendMessage(message);
                        log.debug("消息已发送给用户 {}，会话ID: {}", openId, s.getId());
                    } catch (IOException e) {
                        log.error("向用户 {} 发送消息失败: {}", openId, e.getMessage());
                    }
                }
            }
        } else {
            log.debug("用户 {} 暂无活跃连接，消息无法送达", openId);
        }
    }

    private void broadcastToGroup(String groupId, TextMessage message) {
        List<WebSocketSession> sessions = sessionMap.get(groupId); // 仍沿用原有“用户名即群组”假设
        if (sessions != null && !sessions.isEmpty()) {
            int successCount = 0;
            for (WebSocketSession s : sessions) {
                if (s.isOpen()) {
                    try {
                        s.sendMessage(message);
                        successCount++;
                    } catch (IOException e) {
                        log.error("向群组 {} 广播消息失败: {}", groupId, e.getMessage());
                    }
                }
            }
            log.debug("向群组 {} 广播消息，成功发送: {}/{}", groupId, successCount, sessions.size());
        } else {
            log.debug("群组 {} 不存在或没有活跃连接", groupId);
        }
    }

    private void broadcastToAll(TextMessage message) {
        int totalSessions = 0;
        int successCount = 0;

        for (List<WebSocketSession> sessions : sessionMap.values()) {
            for (WebSocketSession s : sessions) {
                totalSessions++;
                if (s.isOpen()) {
                    try {
                        s.sendMessage(message);
                        successCount++;
                    } catch (IOException e) {
                        log.error("广播消息失败: {}", e.getMessage());
                    }
                }
            }
        }
        log.debug("全局广播完成，成功发送: {}/{}", successCount, totalSessions);
    }

    /* ======================== Payload 工具 ======================== */

    private String createMessagePayload(String from, String content, String type) {
        try {
            Map<String, Object> payload = new HashMap<>();
            payload.put("type", type);
            payload.put("from", from);
            payload.put("content", content);
            payload.put("timestamp", System.currentTimeMillis());
            return OBJECT_MAPPER.writeValueAsString(payload);
        } catch (JsonProcessingException e) {
            log.error("创建消息负载失败: {}", e.getMessage());
            return createErrorPayload("E_JSON", "消息格式错误");
        }
    }

    private String createDialogPayload(String dialogId, String from, String content) {
        try {
            Map<String, Object> payload = new HashMap<>();
            payload.put("type", TYPE_DIALOG_MSG);
            payload.put("dialogId", dialogId);
            payload.put("from", from);
            payload.put("content", content);
            payload.put("timestamp", System.currentTimeMillis());
            return OBJECT_MAPPER.writeValueAsString(payload);
        } catch (JsonProcessingException e) {
            return createErrorPayload("E_JSON", "消息格式错误");
        }
    }

    private String createSystemPayload(String event, String content) {
        try {
            Map<String, Object> payload = new HashMap<>();
            payload.put("type", "SYSTEM");
            payload.put("event", event);
            payload.put("content", content);
            payload.put("timestamp", System.currentTimeMillis());
            return OBJECT_MAPPER.writeValueAsString(payload);
        } catch (JsonProcessingException e) {
            return "{\"error\":\"系统消息格式错误\"}";
        }
    }

    private String createSystemEvent(String event, String content, Map<String, Object> extra) {
        try {
            Map<String, Object> payload = new HashMap<>();
            payload.put("type", "SYSTEM");
            payload.put("event", event);
            payload.put("content", content);
            if (extra != null) payload.putAll(extra);
            payload.put("timestamp", System.currentTimeMillis());
            return OBJECT_MAPPER.writeValueAsString(payload);
        } catch (JsonProcessingException e) {
            return "{\"error\":\"系统消息格式错误\"}";
        }
    }

    private String createErrorPayload(String code, String msg) {
        try {
            Map<String, Object> payload = new HashMap<>();
            payload.put("type", "ERROR");
            payload.put("code", code);
            payload.put("content", msg);
            payload.put("timestamp", System.currentTimeMillis());
            return OBJECT_MAPPER.writeValueAsString(payload);
        } catch (JsonProcessingException e) {
            return "{\"error\":\"错误消息格式错误\"}";
        }
    }

    private static Map<String, Object> mapOf(Object... kv) {
        Map<String, Object> m = new HashMap<>();
        for (int i = 0; i + 1 < kv.length; i += 2) {
            m.put(String.valueOf(kv[i]), kv[i + 1]);
        }
        return m;
    }

    /* ======================== 指标获取 ======================== */

    public int getSessionCountByUser(String openId) {
        List<WebSocketSession> sessions = sessionMap.get(openId);
        return sessions != null ? sessions.size() : 0;
    }

    public int getUserCount() {
        return sessionMap.size();
    }

    public int getTotalSessionCount() {
        return sessionMap.values().stream().mapToInt(List::size).sum();
    }
}
